{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7TX0mtitYYN_"
      },
      "source": [
        "# Retrieval-augmented Thoughts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FgLhY-ICYxNu"
      },
      "outputs": [],
      "source": [
        "# Preparation\n",
        "from openai import OpenAI\n",
        "api_key = \"\" # Your API key\n",
        "openai_client = OpenAI(api_key=api_key)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V8ox9jk5Y_ZC"
      },
      "outputs": [],
      "source": [
        "# Basic Tool Functions\n",
        "import os\n",
        "os.environ[\"GOOGLE_CSE_ID\"] = \"\" # Your Google CSE ID\n",
        "os.environ[\"GOOGLE_API_KEY\"] = \"\" # Your Google API Key\n",
        "\n",
        "from langchain.tools import Tool\n",
        "from langchain_community.utilities import GoogleSearchAPIWrapper\n",
        "def get_search(query:str=\"\", k:int=1): # get the top-k resources with google\n",
        "    search = GoogleSearchAPIWrapper(k=k)\n",
        "    def search_results(query):\n",
        "        return search.results(query, k)\n",
        "    tool = Tool(\n",
        "        name=\"Google Search Snippets\",\n",
        "        description=\"Search Google for recent results.\",\n",
        "        func=search_results,\n",
        "    )\n",
        "    ref_text = tool.run(query)\n",
        "    if 'Result' not in ref_text[0].keys():\n",
        "        return ref_text\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "from langchain_community.document_transformers import Html2TextTransformer\n",
        "from langchain_community.document_loaders import AsyncHtmlLoader\n",
        "def get_page_content(link:str):\n",
        "    loader = AsyncHtmlLoader([link])\n",
        "    docs = loader.load()\n",
        "    html2text = Html2TextTransformer()\n",
        "    docs_transformed = html2text.transform_documents(docs)\n",
        "    if len(docs_transformed) > 0:\n",
        "        return docs_transformed[0].page_content\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "import tiktoken\n",
        "def num_tokens_from_string(string: str, encoding_name: str = \"cl100k_base\") -> int:\n",
        "    \"\"\"Returns the number of tokens in a text string.\"\"\"\n",
        "    encoding = tiktoken.get_encoding(encoding_name)\n",
        "    num_tokens = len(encoding.encode(string))\n",
        "    return num_tokens\n",
        "\n",
        "def chunk_text_by_sentence(text, chunk_size=2048):\n",
        "    \"\"\"Chunk the $text into sentences with less than 2k tokens.\"\"\"\n",
        "    sentences = text.split('. ')\n",
        "    chunked_text = []\n",
        "    curr_chunk = []\n",
        "    # 逐句添加文本片段，确保每个段落都小于2k个token\n",
        "    for sentence in sentences:\n",
        "        if num_tokens_from_string(\". \".join(curr_chunk)) + num_tokens_from_string(sentence) + 2 <= chunk_size:\n",
        "            curr_chunk.append(sentence)\n",
        "        else:\n",
        "            chunked_text.append(\". \".join(curr_chunk))\n",
        "            curr_chunk = [sentence]\n",
        "    # 添加最后一个片段\n",
        "    if curr_chunk:\n",
        "        chunked_text.append(\". \".join(curr_chunk))\n",
        "    return chunked_text[0]\n",
        "\n",
        "def chunk_text_front(text, chunk_size = 2048):\n",
        "    '''\n",
        "    get the first `trunk_size` token of text\n",
        "    '''\n",
        "    chunked_text = \"\"\n",
        "    tokens = num_tokens_from_string(text)\n",
        "    if tokens < chunk_size:\n",
        "        return text\n",
        "    else:\n",
        "        ratio = float(chunk_size) / tokens\n",
        "        char_num = int(len(text) * ratio)\n",
        "        return text[:char_num]\n",
        "\n",
        "def chunk_texts(text, chunk_size = 2048):\n",
        "    '''\n",
        "    trunk the text into n parts, return a list of text\n",
        "    [text, text, text]\n",
        "    '''\n",
        "    tokens = num_tokens_from_string(text)\n",
        "    if tokens < chunk_size:\n",
        "        return [text]\n",
        "    else:\n",
        "        texts = []\n",
        "        n = int(tokens/chunk_size) + 1\n",
        "        # 计算每个部分的长度\n",
        "        part_length = len(text) // n\n",
        "        # 如果不能整除，则最后一个部分会包含额外的字符\n",
        "        extra = len(text) % n\n",
        "        parts = []\n",
        "        start = 0\n",
        "\n",
        "        for i in range(n):\n",
        "            # 对于前extra个部分，每个部分多分配一个字符\n",
        "            end = start + part_length + (1 if i < extra else 0)\n",
        "            parts.append(text[start:end])\n",
        "            start = end\n",
        "        return parts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4HBQ5onEn_m-"
      },
      "outputs": [],
      "source": [
        "# RAT Pipeline\n",
        "\n",
        "from datetime import datetime\n",
        "chatgpt_system_prompt = f'''\n",
        "You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture.\n",
        "Knowledge cutoff: 2023-04\n",
        "Current date: {datetime.now().strftime('%Y-%m-%d')}\n",
        "'''\n",
        "\n",
        "def get_draft(question):\n",
        "    # Getting the draft answer\n",
        "    draft_prompt = '''\n",
        "IMPORTANT:\n",
        "Try to answer this question/instruction with step-by-step thoughts and make the answer more structural.\n",
        "Use `\\n\\n` to split the answer into several paragraphs.\n",
        "Just respond to the instruction directly. DO NOT add additional explanations or introducement in the answer unless you are asked to.\n",
        "'''\n",
        "    draft = openai_client.chat.completions.create(\n",
        "        model=\"gpt-3.5-turbo\",\n",
        "        messages=[\n",
        "            {\n",
        "                \"role\": \"system\",\n",
        "                \"content\": chatgpt_system_prompt\n",
        "            },\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"content\": f\"{question}\" + draft_prompt\n",
        "            }\n",
        "        ],\n",
        "        temperature = 1.0\n",
        "    ).choices[0].message.content\n",
        "    return draft\n",
        "\n",
        "def split_draft(draft, split_char = '\\n\\n'):\n",
        "    # 将draft切分为多个段落\n",
        "    # split_char: '\\n\\n'\n",
        "    draft_paragraphs = draft.split(split_char)\n",
        "    # print(f\"The draft answer has {len(draft_paragraphs)}\")\n",
        "    return draft_paragraphs\n",
        "\n",
        "def get_query(question, answer):\n",
        "    query_prompt = '''\n",
        "I want to verify the content correctness of the given question, especially the last sentences.\n",
        "Please summarize the content with the corresponding question.\n",
        "This summarization will be used as a query to search with Bing search engine.\n",
        "The query should be short but need to be specific to promise Bing can find related knowledge or pages.\n",
        "You can also use search syntax to make the query short and clear enough for the search engine to find relevant language data.\n",
        "Try to make the query as relevant as possible to the last few sentences in the content.\n",
        "**IMPORTANT**\n",
        "Just output the query directly. DO NOT add additional explanations or introducement in the answer unless you are asked to.\n",
        "'''\n",
        "    query = openai_client.chat.completions.create(\n",
        "        model=\"gpt-3.5-turbo\",\n",
        "        messages=[\n",
        "            {\n",
        "                \"role\": \"system\",\n",
        "                \"content\": chatgpt_system_prompt\n",
        "            },\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"content\": f\"##Question: {question}\\n\\n##Content: {answer}\\n\\n##Instruction: {query_prompt}\"\n",
        "            }\n",
        "        ],\n",
        "        temperature = 1.0\n",
        "    ).choices[0].message.content\n",
        "    return query\n",
        "\n",
        "def get_content(query):\n",
        "    res = get_search(query, 1)\n",
        "    if not res:\n",
        "        print(\">>> No good Google Search Result was found\")\n",
        "        return None\n",
        "    search_results = res[0]\n",
        "    link = search_results['link'] # title, snippet\n",
        "    res = get_page_content(link)\n",
        "    if not res:\n",
        "        print(f\">>> No content was found in {link}\")\n",
        "        return None\n",
        "    retrieved_text = res\n",
        "    trunked_texts = chunk_texts(retrieved_text, 1500)\n",
        "    trunked_texts = [trunked_text.replace('\\n', \" \") for trunked_text in trunked_texts]\n",
        "    return trunked_texts\n",
        "\n",
        "def get_revise_answer(question, answer, content):\n",
        "    revise_prompt = '''\n",
        "I want to revise the answer according to retrieved related text of the question in WIKI pages.\n",
        "You need to check whether the answer is correct.\n",
        "If you find some errors in the answer, revise the answer to make it better.\n",
        "If you find some necessary details are ignored, add it to make the answer more plausible according to the related text.\n",
        "If you find the answer is right and do not need to add more details, just output the original answer directly.\n",
        "**IMPORTANT**\n",
        "Try to keep the structure (multiple paragraphs with its subtitles) in the revised answer and make it more structual for understanding.\n",
        "Split the paragraphs with `\\n\\n` characters.\n",
        "Just output the revised answer directly. DO NOT add additional explanations or annoucement in the revised answer unless you are asked to.\n",
        "'''\n",
        "    revised_answer = openai_client.chat.completions.create(\n",
        "        model=\"gpt-3.5-turbo\",\n",
        "        messages=[\n",
        "                {\n",
        "                    \"role\": \"system\",\n",
        "                    \"content\": chatgpt_system_prompt\n",
        "                },\n",
        "                {\n",
        "                    \"role\": \"user\",\n",
        "                    \"content\": f\"##Existing Text in Wiki Web: {content}\\n\\n##Question: {question}\\n\\n##Answer: {answer}\\n\\n##Instruction: {revise_prompt}\"\n",
        "                }\n",
        "            ],\n",
        "            temperature = 1.0\n",
        "    ).choices[0].message.content\n",
        "    return revised_answer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yW1yp-BQQUxp"
      },
      "outputs": [],
      "source": [
        "def get_query_wrapper(q, question, answer):\n",
        "    result = get_query(question, answer)\n",
        "    q.put(result)  # 将结果放入队列\n",
        "\n",
        "def get_content_wrapper(q, query):\n",
        "    result = get_content(query)\n",
        "    q.put(result)  # 将结果放入队列\n",
        "\n",
        "def get_revise_answer_wrapper(q, question, answer, content):\n",
        "    result = get_revise_answer(question, answer, content)\n",
        "    q.put(result)\n",
        "\n",
        "from multiprocessing import Process, Queue\n",
        "def run_with_timeout(func, timeout, *args, **kwargs):\n",
        "    q = Queue()  # 创建一个Queue对象用于进程间通信\n",
        "    # 创建一个进程来执行传入的函数，将Queue和其他*args、**kwargs作为参数传递\n",
        "    p = Process(target=func, args=(q, *args), kwargs=kwargs)\n",
        "    p.start()\n",
        "    # 等待进程完成或超时\n",
        "    p.join(timeout)\n",
        "    if p.is_alive():\n",
        "        print(f\"{datetime.now()} [INFO] 函数{str(func)}执行已超时({timeout}s)，正在终止进程...\")\n",
        "        p.terminate()  # 终止进程\n",
        "        p.join()  # 确保进程已经终止\n",
        "        result = None  # 超时情况下，我们没有结果\n",
        "    else:\n",
        "        print(f\"{datetime.now()} [INFO] 函数{str(func)}执行成功完成\")\n",
        "        result = q.get()  # 从队列中获取结果\n",
        "    return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from difflib import unified_diff\n",
        "from IPython.display import display, HTML\n",
        "\n",
        "def generate_diff_html(text1, text2):\n",
        "    diff = unified_diff(text1.splitlines(keepends=True),\n",
        "                        text2.splitlines(keepends=True),\n",
        "                        fromfile='text1', tofile='text2')\n",
        "\n",
        "    diff_html = \"\"\n",
        "    for line in diff:\n",
        "        if line.startswith('+'):\n",
        "            diff_html += f\"<div style='color:green;'>{line.rstrip()}</div>\"\n",
        "        elif line.startswith('-'):\n",
        "            diff_html += f\"<div style='color:red;'>{line.rstrip()}</div>\"\n",
        "        elif line.startswith('@'):\n",
        "            diff_html += f\"<div style='color:blue;'>{line.rstrip()}</div>\"\n",
        "        else:\n",
        "            diff_html += f\"{line.rstrip()}<br>\"\n",
        "    return diff_html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jE3Y78zfoLRc"
      },
      "outputs": [],
      "source": [
        "# RAT Function\n",
        "newline_char = '\\n'\n",
        "\n",
        "def rat(question):\n",
        "    print(f\"{datetime.now()} [INFO] 获取草稿...\")\n",
        "    draft = get_draft(question)\n",
        "    print(f\"{datetime.now()} [INFO] 返回草稿\")\n",
        "    print(f\"##################### DRAFT #######################\")\n",
        "    print(draft)\n",
        "    print(f\"#####################  END  #######################\")\n",
        "\n",
        "    print(f\"{datetime.now()} [INFO] 处理草稿...\")\n",
        "    draft_paragraphs = split_draft(draft)\n",
        "    print(f\"{datetime.now()} [INFO] 草稿被切分为{len(draft_paragraphs)}部分\")\n",
        "    answer = \"\"\n",
        "    for i, p in enumerate(draft_paragraphs):\n",
        "        print(str(i)*80)\n",
        "        print(f\"{datetime.now()} [INFO] 修改第{i+1}/{len(draft_paragraphs)}部分...\")\n",
        "        answer = answer + '\\n\\n' + p\n",
        "        # print(f\"[{i}/{len(draft_paragraphs)}] Original Answer:\\n{answer.replace(newline_char, ' ')}\")\n",
        "\n",
        "        # query = get_query(question, answer)\n",
        "        print(f\"{datetime.now()} [INFO] 生成对应Query...\")\n",
        "        res = run_with_timeout(get_query_wrapper, 3, question, answer)\n",
        "        if not res:\n",
        "            print(f\"{datetime.now()} [INFO] 跳过后续步骤...\")\n",
        "            continue\n",
        "        else:\n",
        "            query = res\n",
        "        print(f\">>> {i}/{len(draft_paragraphs)} Query: {query.replace(newline_char, ' ')}\")\n",
        "\n",
        "        print(f\"{datetime.now()} [INFO] 获取网页内容...\")\n",
        "        # content = get_content(query)\n",
        "        res = run_with_timeout(get_content_wrapper, 5, query)\n",
        "        if not res:\n",
        "            print(f\"{datetime.now()} [INFO] 跳过后续步骤...\")\n",
        "            continue\n",
        "        else:\n",
        "            content = res\n",
        "\n",
        "        for j, c in enumerate(content):\n",
        "            if  j > 2:\n",
        "                break\n",
        "            print(f\"{datetime.now()} [INFO] 根据网页内容修改对应答案...[{j}/{min(len(content),3)}]\")\n",
        "            # answer = get_revise_answer(question, answer, c)\n",
        "            res = run_with_timeout(get_revise_answer_wrapper, 10, question, answer, c)\n",
        "            if not res:\n",
        "                print(f\"{datetime.now()} [INFO] 跳过后续步骤...\")\n",
        "                continue\n",
        "            else:\n",
        "                diff_html = generate_diff_html(answer, res)\n",
        "                display(HTML(diff_html))\n",
        "                answer = res\n",
        "            print(f\"{datetime.now()} [INFO] 答案修改完成[{j}/{min(len(content),3)}]\")\n",
        "        # print(f\"[{i}/{len(draft_paragraphs)}] REVISED ANSWER:\\n {answer.replace(newline_char, ' ')}\")\n",
        "        # print()\n",
        "    return draft, answer\n",
        "    # return answer\n",
        "\n",
        "# draft, answer = rat(\"介绍金庸的生平(写3段)\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vr7TsTHUes7k",
        "outputId": "41bff12a-c2d1-45b0-d15f-eccdfe572b7e"
      },
      "outputs": [],
      "source": [
        "draft, answer = rat(\"Introduce Jin-Yong's Life.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "diff_html = generate_diff_html(draft, answer)\n",
        "display(HTML(diff_html))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c71QsU8SYsC2"
      },
      "source": [
        "# Gradio Demo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BGeb6Q9FZnq8"
      },
      "outputs": [],
      "source": [
        "def predict(message, history):\n",
        "    # rat function?\n",
        "    draft, answer = rat(message)\n",
        "    return answer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "YnFyowIFYtzE",
        "outputId": "3d2b8b2c-b118-4a5b-bf90-3a50f4f7fb56"
      },
      "outputs": [],
      "source": [
        "# Fast ChatBot demo\n",
        "import gradio as gr\n",
        "gr.ChatInterface(predict).launch(share=True, debug = True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
